"""
    DH 秘钥协商
"""
from common.math import quick_pow

class DH:
    def __init__(self,sender_pub_hex_arg:str,receiver_pub_hex_arg:str):
        self.p = int(sender_pub_hex_arg,16)
        self.g = int(receiver_pub_hex_arg,16)
        self._dh = _DH(self.p,self.g)

    def generate_pub_key(self,secret_hex_arg:str,to_hex_str:bool=True):
        secret = int(secret_hex_arg,16)
        pub = self._dh.generate_pub_key(secret)
        return hex(pub) if to_hex_str else pub

    def calcu_shared_key(self,self_secret:str,the_other_side_pub:str,to_hex_str:bool=True):
        secret = int(self_secret,16)
        A = int(the_other_side_pub,16)
        s = self._dh.calcu_shared_key(secret,A)
        return hex(s) if to_hex_str else s


class _DH:
    def __init__(self,p:int,g:int):
        self.p = p
        self.g = g

    def generate_pub_key(self,secret:int):
        pub = quick_pow(self.g,secret,self.p)
        return pub 
        
    def calcu_shared_key(self,secret,A):
        shared_key = quick_pow(A,secret,self.p)
        return shared_key



if __name__=="__main__":
    alice = (1234,1111)
    bob = (4321,2222)
    p,g = alice[0],bob[0]
    a,b = alice[1],bob[1]
    dh = DH(hex(p),hex(g))
    A = dh.generate_pub_key(hex(a),False)
    B = dh.generate_pub_key(hex(b),False)
    print(dh.calcu_shared_key(hex(a),hex(B)))
    print(dh.calcu_shared_key(hex(b),hex(A)))
    


    